import multiprocessing
from importlib import import_module
from datetime import datetime
import numpy as np

from python.inequalities import Inequality
import python.version as version
from python.propagation_data import (
    LWEInstance,
    PropagationData,
    PropagationDataStep,
)
from python.solve import solve
from python.guass_elimate import solve_equation
from python.helpers import (
    check_inequalities,
    evol_bino,
    bino,
    IneqType,
    print_v
)


# python_kyber = import_module(f"python_kyber{version.KYBER_VERSION}")


def propagate(
        key,
        graph,
        count_steps=1,
        step_size=1,
        thread_count=None,
        propagation_data=None,
):
    if thread_count is None:
        thread_count = multiprocessing.cpu_count()
    success = False
    step_size *= 2
    results_list = graph.get_results(thread_count)
    if propagation_data:
        step_data = PropagationDataStep.new(results_list, key, 0)
        propagation_data.add_step(0, step_data)
    # 没有propagation是否应当有一个else进行报错
    print_v(f"Using {thread_count} threads.\n")
    propagation_data.set_start(datetime.now())
    for step in range(0, 2 * count_steps, step_size):
        with open('output.txt', 'a') as file:
            file.write(f"----Propagation step {step // 2}----")
            print(f"changheng----Propagation step {step // 2}----")
        graph.propagate(step_size, thread_count)
        print_v("Fetching results..")
        results_list = graph.get_results(thread_count)
        # results = {i: results_list[i] for i in range(len(key))}
        step_data = PropagationDataStep.new(results_list, key, step + step_size)
        if propagation_data:
            propagation_data.add_step(step + step_size, step_data)
        print_v(f"Average entropy is {step_data.avg_entropy}.")
        print_v(f"Maximal entropy is {step_data.max_entropy}.")
        print_v("")
        with open('output.txt', 'a') as file:
            file.write(f"{step_data.recovered_coefficients} coefficients recovered")
            file.write("\n")
            print(f"changyesheng{step_data.recovered_coefficients} coefficients recovered")
        if step_data.guessed_key == key:
            print_v("Found correct key.")
            success = True
            break
        if step_data.recovered_coefficients >= version.KYBER_k * version.KYBER_N:
            print_v("Found enough correct coefficients")
            success = True
            break
    propagation_data.set_end(datetime.now())
    if success:
        with open('output.txt', 'a') as file:
            file.write("BP alone: Success!")
            file.write("\n")
            print("BP alone: Success!")
    else:
        with open('output.txt', 'a') as file:
            file.write("BP alone: Failure!")
            file.write("\n")
            print("BP alone: Failure!")
    return success


def create_graph_inequalities(ineqs, sum_dist, dist, thresholds, sweep_efforts):
    check_bp = import_module(f"check_bp{version.KYBER_VERSION}")
    print_v("Building check graph..")
    g = check_bp.CheckGraph()
    g.add_var_nodes(sum_dist, dist)
    total = len(ineqs)
    corrects = 0
    for index, ineq in enumerate(ineqs):
        if ineq.p_correct < 1.0:
            # 概率p正确不等式
            g.add_inequality_prob(
                f"Line {index}",
                index,
                ineq.coefficients,
                ineq.b,
                # ---------------------------------
                ineq.is_double,
                # -----------------------------------
                ineq.sign == IneqType.LE,
                ineq.p_correct,
            )
        else:
            # 概率1正确不等式
            g.add_inequality(
                f"Line {index}",
                index,
                ineq.coefficients,
                ineq.b,
                # ---------------------------------
                ineq.is_double,
                # -----------------------------------
                ineq.sign == IneqType.LE,
                thresholds,
                sweep_efforts,
            )
            corrects += 1
        print_v(f"{index}/{total}\t\t", end="\r")
    print_v("                                          ")
    print_v(f"Created {total} inequalities, {corrects} are certainly correct, {total - corrects} might be incorrect.")
    print_v("Initializing graph..")
    g.ini()
    return g


def load_my_inequalities(path, no_ineqs, bilateral, narrowing_path=None, group=0):
    inequalities = []
    # read coeff matrix a
    a_file = open(path + ".txt", 'r')
    # read lower bound, value, upper bound
    lvu_file = open(path + "_lvu.txt", 'r')

    a_lines = a_file.readlines()[group * no_ineqs:]
    lvu_lines = lvu_file.readlines()[group * no_ineqs:]

    narrowing_eqs = {}  # a dictionary storing narrowing information, keys: eq_indexes, value: se

    # read which inequality narrows and how much it narrows
    if narrowing_path is not None:
        narrowing_file = open(narrowing_path, 'r')
        for line in narrowing_file:
            line = line.split(', ')
            if line[0] == 'narrowing':
                eq_index = int(line[1])
                se = int(line[2])
                if eq_index in narrowing_eqs.keys():
                    narrowing_eqs[eq_index] += se
                else:
                    narrowing_eqs[eq_index] = se
        narrowing_file.close()

    l, v, u, dropped = [], [], [], []  # 用于记录丢掉的proximity大的方程
    i, d = 0, 0
    for line in lvu_lines:  # 第i组方程的上下界
        line = list(map(int, line.split()))

        if i in narrowing_eqs:
            l.append(line[0] + narrowing_eqs[i])
            u.append(line[2] - narrowing_eqs[i])

        else:
            if min(line[2] - line[1], line[1] - line[0]) > 1000:  # 这里设置成 > 1000那就是不丢了
                dropped.append(d)
                d += 1
                continue
            else:
                l.append(line[0])
                u.append(line[2])
        v.append(line[1])
        i += 1
        d += 1

        # 如果已经处理完了所需的不等式，提前退出循环
        if i >= no_ineqs:
            break

    i = 0
    if bilateral:  # 双边不等式求解
        for index, line in enumerate(a_lines):
            if index in dropped:
                continue
            line = [int(x) for x in line.split()]

            if i in narrowing_eqs.keys():
                inequalities.append(
                    Inequality(line, True, IneqType.DB, [l[i], u[i], v[i]], False, 0.9965)
                )
                # 双边不等式，保存[下界，上界，真实值]. 以概率  0.9 成立
            else:
                inequalities.append(
                    Inequality(line, True, IneqType.DB, [l[i], u[i], v[i]], True, 1)
                )
                # 双边不等式，保存[下界，上界，真实值]. 以概率 1 成立
            i += 1

            if i >= no_ineqs:
                break

    else:  # 按照单边不等式求解
        for index, line in enumerate(a_lines):
            if index in dropped:
                continue
            line = [int(x) for x in line.split()]

            if i in narrowing_eqs:
                inequalities.append(Inequality(line, False, IneqType.GE, [l[i], 0, v[i]], False, 0.9))
                inequalities.append(Inequality(line, False, IneqType.LE, [u[i], 0, v[i]], False, 0.9))
            else:
                inequalities.append(Inequality(line, False, IneqType.GE, [l[i], 0, v[i]], True, 1))
                inequalities.append(Inequality(line, False, IneqType.LE, [u[i], 0, v[i]], True, 1))

            i += 1

            if i >= no_ineqs:
                break

    # load right key
    es_file = open(path + "_es.txt", 'r')
    key = [int(line) for line in es_file]

    assert check_inequalities(key, inequalities)
    print_v("                                          ")
    print_v("Number of inequalities: ", no_ineqs)
    print_v("Finished loading my inequalities.")

    # load lwe_A
    lwe_A_file = open(path + "_pkA.txt", 'r')
    m = version.KYBER_k * version.KYBER_N
    lwe_A = np.zeros(shape=(m, m), dtype=np.int64)
    for i, line in enumerate(lwe_A_file):
        lwe_A[i] = [int(x) for x in line.split()]

    # load lwe_t
    lwe_t_file = open(path + "_pkt.txt", 'r')
    lwe_t = [int(line) for line in lwe_t_file]

    propagation_data = PropagationData.new(
        key=key,
        inequalities=inequalities,
        filtered_cts=None,
        max_delta_v=None,
        lwe_instance=LWEInstance(lwe_A.tolist(), lwe_t)
    )

    lwe_A_file.close()
    lwe_t_file.close()
    es_file.close()
    a_file.close()
    lvu_file.close()

    return propagation_data


def run_with_inequalities(
        propagation_data,
        steps,
        step_size,
        block_size,
        run_reduction,
        max_beta,
        add_fplll,
        use_best_step,
        max_enum,
        thread_count,
        thresholds,
        sweep_efforts,
):
    if len(propagation_data.steps) == 0:
        g = create_graph_inequalities(
            propagation_data.inequalities,
            evol_bino(version.ETA1),
            bino(version.ETA1),
            thresholds,
            sweep_efforts
        )
        success_bp = propagate(
            propagation_data.key,
            g,
            steps,
            step_size,
            propagation_data=propagation_data,
            thread_count=thread_count
        )
        # 这里返回一个bool
    elif propagation_data.steps:
        steps = sorted(propagation_data.steps.items(), key=lambda x: x[0])
        success_bp = steps[-1][1].recovered_coefficients
        # 这里返回一个int 不一样也ok吗
    else:
        success_bp = None
    print_v("BP alone " + "succeeded" if success_bp else "did not succeed" + ".")
    # 找到正确系数最多的步骤的序号
    step_idx = (
        max(propagation_data.steps, key=lambda x: propagation_data.steps[x].correct_coefficients)
        if use_best_step
        else -1
    )
    # 按步骤序号排序
    step_list = sorted(propagation_data.steps.items(), key=lambda x: x[0])
    pos = 0
    if step_idx == -1:
        step_idx = step_list[-1][0]
    else:
        for idx, _ in reversed(step_list):
            if idx == step_idx:
                break
            pos += 1
    with open('output.txt', 'a') as file:
        file.write(f"Using step {step_idx} which is on position {pos}," +
                   f"has {propagation_data.steps[step_idx].correct_coefficients} correct coefficients" +
                   f"and {propagation_data.steps[step_idx].recovered_coefficients} recovered coefficients.\n")
        print(f"changheng Using step {step_idx} which is on position {pos}," +
              f"has {propagation_data.steps[step_idx].correct_coefficients} correct coefficients" +
              f"and {propagation_data.steps[step_idx].recovered_coefficients} recovered coefficients.", end='')
    propagation_data.recovered_coefficients = propagation_data.steps[step_idx].recovered_coefficients
    if success_bp:
        success = solve_equation(
            propagation_data,
            step_idx,
            pos
        )
    else:
        success = solve(
            propagation_data,
            block_size=block_size,
            max_beta=max_beta,
            perform=run_reduction,
            add_fplll=add_fplll,
            step=step_idx,
            max_enum=max_enum,
            step_rank=pos,
        )
    return success, propagation_data
